import random
import cv2
import albumentations as albu
from albumentations.augmentations import functional as F
from albumentations.augmentations.bbox_utils import denormalize_bbox, normalize_bbox

__all__ = ['init_test_kp_albu', 'init_train_kp_albu', 'init_train_seg_albu', 'init_test_seg_albu', 'init_train_kp_albu_simple']



class PadScale(albu.DualTransform):
    """Pad side of the image / max if side is less than desired number.

    Args:
        min_height (int): minimal result image height.
        min_width (int): minimal result image width.
        border_mode (OpenCV flag): OpenCV border mode.
        value (int, float, list of int, lisft of float): padding value if border_mode is cv2.BORDER_CONSTANT.
        mask_value (int, float,
                    list of int,
                    lisft of float): padding value for mask if border_mode is cv2.BORDER_CONSTANT.
        p (float): probability of applying the transform. Default: 1.0.

    Targets:
        image, mask, bbox, keypoints

    Image types:
        uint8, float32
    """

    def __init__(
        self,
        max_scale=1.5,
        border_mode=cv2.BORDER_REFLECT_101,
        value=None,
        mask_value=None,
        always_apply=False,
        p=1.0,
    ):
        super(PadScale, self).__init__(always_apply, p)
        self.max_scale = max_scale
        self.border_mode = border_mode
        self.value = value
        self.mask_value = mask_value

    def update_params(self, params, **kwargs):
        params = super(PadScale, self).update_params(params, **kwargs)
        rows = params["rows"]
        cols = params["cols"]

        ratio = random.uniform(1, self.max_scale)
        y = random.uniform(0, rows * ratio - rows)
        x = random.uniform(0, cols * ratio - cols)

        if y > 0:
            h_pad_top = int(y / 2.0)
            h_pad_bottom = h_pad_top
        else:
            h_pad_top = 0
            h_pad_bottom = 0

        if cols > 0:
            w_pad_left = int(x / 2.0)
            w_pad_right = w_pad_left
        else:
            w_pad_left = 0
            w_pad_right = 0

        params.update(
            {"pad_top": h_pad_top, "pad_bottom": h_pad_bottom, "pad_left": w_pad_left, "pad_right": w_pad_right}
        )
        return params

    def apply(self, img, pad_top=0, pad_bottom=0, pad_left=0, pad_right=0, **params):
        return F.pad_with_params(
            img, pad_top, pad_bottom, pad_left, pad_right, border_mode=self.border_mode, value=self.value
        )

    def apply_to_mask(self, img, pad_top=0, pad_bottom=0, pad_left=0, pad_right=0, **params):
        return F.pad_with_params(
            img, pad_top, pad_bottom, pad_left, pad_right, border_mode=self.border_mode, value=self.mask_value
        )

    def apply_to_bbox(self, bbox, pad_top=0, pad_bottom=0, pad_left=0, pad_right=0, rows=0, cols=0, **params):
        x_min, y_min, x_max, y_max = denormalize_bbox(bbox, rows, cols)
        bbox = x_min + pad_left, y_min + pad_top, x_max + pad_left, y_max + pad_top
        return normalize_bbox(bbox, rows + pad_top + pad_bottom, cols + pad_left + pad_right)

    # skipcq: PYL-W0613
    def apply_to_keypoint(self, keypoint, pad_top=0, pad_bottom=0, pad_left=0, pad_right=0, **params):
        x, y, angle, scale = keypoint
        return x + pad_left, y + pad_top, angle, scale

    def get_transform_init_args_names(self):
        return ("min_height", "min_width", "border_mode", "value", "mask_value")



# training keypoint albu
def init_train_kp_albu(height, width):
    transformations = albu.Compose([
        albu.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=0.5),

        # albu.RandomScale(scale_limit=0.10, interpolation=2, p=0.5),
        # albu.PadIfNeeded(min_height=height, min_width=width, border_mode=cv2.BORDER_CONSTANT, value=0, p=1.0),
        # albu.RandomCrop(height, width, p=1.0),
        #
        albu.HorizontalFlip(p=0.5),
        albu.OneOf(
            [albu.Rotate(limit=5, interpolation=1, border_mode=cv2.BORDER_CONSTANT, value=0, p=0.5),
             #albu.IAAPiecewiseAffine(scale=(0.03, 0.05), p=0.2),
             #albu.IAAPerspective(scale=(0.01, 0.05), p=0.2)
             ],p=0.5
        ),

        #
        albu.OneOf([  # albu.Blur(blur_limit=3, p=0.5),
            albu.GaussianBlur(blur_limit=3, p=0.5),
            albu.MotionBlur(blur_limit=3, p=0.5),
            #albu.ISONoise(color_shift=(0.01, 0.05), intensity=(0.1, 0.5), p=0.5)
            ], p=0.5),


        albu.OneOf([albu.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.1, p=0.3),
                    albu.RandomGamma(gamma_limit=(80, 120), p=0.3),
                    # albu.RGBShift(r_shift_limit=20, g_shift_limit=20, b_shift_limit=20, p=0.5),
                    # albu.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5)
                    ], p=0.5
                   ),

        # albu.CoarseDropout(max_holes=8, max_height=16, max_width=16, p=0.1)


        albu.Resize(height=height, width=width, interpolation=cv2.INTER_CUBIC, p=1.0),


    ])

    target = {}
    target['image'] = 'image'
    target['mask'] = 'mask'
    target['keypoints'] = 'keypoints'

    return albu.Compose(transformations, p=1.0,
                        keypoint_params=albu.KeypointParams(format='xy', remove_invisible=False),
                        additional_targets=target)

# training keypoint albu
def init_train_kp_albu_simple(height, width):
    transformations = albu.Compose([
        #albu.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=0.5),

        # albu.RandomScale(scale_limit=0.10, interpolation=2, p=0.5),
        # albu.PadIfNeeded(min_height=height, min_width=width, border_mode=cv2.BORDER_CONSTANT, value=0, p=1.0),
        # albu.RandomCrop(height, width, p=1.0),
        #

        albu.OneOf([albu.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.1, p=0.5),
                    albu.RandomGamma(gamma_limit=(80, 120), p=0.5),
                    albu.ChannelShuffle(p=0.5)
                    ], p=1.0
                   ),

        PadScale(max_scale=1.5, border_mode=cv2.BORDER_CONSTANT, value=0.5, p=0.5),

        albu.HorizontalFlip(p=0.5),

        albu.Resize(height=height, width=width, interpolation=cv2.INTER_CUBIC, p=1.0),
    ])

    target = {}
    target['image'] = 'image'
    target['mask'] = 'mask'
    target['keypoints'] = 'keypoints'

    return albu.Compose(transformations, p=1.0,
                        keypoint_params=albu.KeypointParams(format='xy', remove_invisible=False),
                        additional_targets=target)

# training segment albu
def init_train_seg_albu(height, width):
    transformations = albu.Compose([
        #albu.CLAHE(clip_limit=0.4, tile_grid_size=(8,8), p=0.5),
        albu.Resize(height=height, width=width, interpolation=cv2.INTER_CUBIC, p=1.0),
        albu.RandomScale(scale_limit=0.25, interpolation=2, p=1.0),
        albu.PadIfNeeded(min_height=height, min_width=width, border_mode=cv2.BORDER_CONSTANT, value=0, p=1.0),
        albu.RandomCrop(height, width, p=1.0),

        # albu.HorizontalFlip(p=0.5),
        # albu.Rotate(limit=25, interpolation=1, border_mode=cv2.BORDER_CONSTANT, value=0, p=0.5),

        albu.OneOf([  # albu.Blur(blur_limit=3, p=0.5),
            albu.GaussianBlur(blur_limit=3, p=0.3),
            albu.MotionBlur(blur_limit=3, p=0.3),
            #albu.ISONoise(color_shift=(0.01, 0.05), intensity=(0.1, 0.5), p=0.3)
        ], p=0.5),


        albu.OneOf([albu.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=0.3),
                    albu.RandomGamma(gamma_limit=(80, 120), p=0.3),
                    albu.InvertImg(p=0.3    )
                    # albu.RGBShift(r_shift_limit=20, g_shift_limit=20, b_shift_limit=20, p=0.5),
                    # albu.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5)
                    ], p=0.5
                   ),
        #Non-rigid transformations
        albu.OneOf([albu.ElasticTransform(alpha=120* 0.1, sigma=120 * 0.05, alpha_affine=120 * 0.03, p=0.2),
                    albu.OpticalDistortion(distort_limit=0.05, shift_limit=0.05, p=0.2)
                    ], p=0.5)
    ])

    # target = {}
    # target['image'] = 'image'
    # target['mask'] = 'mask'
    #
    # return albu.Compose(transformations, p=1.0,
    #                     additional_targets=target)
    return transformations


def init_test_kp_albu(height, width):
    transformations = albu.Compose([
        #albu.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=1.0),
        albu.Resize(height=height, width=width, interpolation=cv2.INTER_CUBIC, p=1.0),
    ])

    target = {}
    target['image'] = 'image'
    target['mask'] = 'mask'
    target['keypoints'] = 'keypoints'

    return albu.Compose(transformations, p=1.0,
                        keypoint_params=albu.KeypointParams(format='xy', remove_invisible=False),
                        additional_targets=target)


def init_test_seg_albu(height, width):
    transformations = albu.Compose([
        albu.CLAHE(clip_limit=0.4, tile_grid_size=(32, 32), p=1.0),
        albu.Resize(height=height, width=width, interpolation=cv2.INTER_CUBIC, p=1.0),
    ])

    # target = {}
    # target['image'] = 'image'
    # target['mask'] = 'mask'
    #
    # return albu.Compose(transformations, p=1.0,
    #                     additional_targets=target)
    return transformations
